import depthai as dai
import cv2
import numpy as np

# 创建 DepthAI 管道
pipeline = dai.Pipeline()

# 定义摄像头节点
cam_rgb = pipeline.create(dai.node.ColorCamera)
cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_12_MP)  # 设置分辨率为 4056x3040
cam_rgb.setIspScale(1, 3)  # 下采样以达到 1352x1013，接近 1670x1114
cam_rgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.RGB)

# 创建深度节点
mono_left = pipeline.create(dai.node.MonoCamera)
mono_left.setBoardSocket(dai.CameraBoardSocket.LEFT)
mono_left.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)

mono_right = pipeline.create(dai.node.MonoCamera)
mono_right.setBoardSocket(dai.CameraBoardSocket.RIGHT)
mono_right.setResolution(dai.MonoCameraProperties.SensorResolution.THE_400_P)

stereo = pipeline.create(dai.node.StereoDepth)
stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)
mono_left.out.link(stereo.left)
mono_right.out.link(stereo.right)

# 创建输出节点
xout_rgb = pipeline.create(dai.node.XLinkOut)
xout_rgb.setStreamName("rgb")

xout_depth = pipeline.create(dai.node.XLinkOut)
xout_depth.setStreamName("depth")

# 链接节点
cam_rgb.isp.link(xout_rgb.input)
stereo.depth.link(xout_depth.input)

# 启动管道
with dai.Device(pipeline) as device:
    q_rgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
    q_depth = device.getOutputQueue(name="depth", maxSize=4, blocking=False)
    count = 0

    while True:
        in_rgb = q_rgb.get()
        in_depth = q_depth.get()

        frame = in_rgb.getCvFrame()
        depth_frame = in_depth.getFrame()

        # 调整图像大小到 1670x1114
        resized_frame = cv2.resize(frame, (1670, 1114))
        resized_depth_frame = cv2.resize(depth_frame, (1670, 1114))

        # 显示图像
        cv2.imshow("RGB Frame", resized_frame)

        # 归一化并显示深度图
        max_depth = np.amax(resized_depth_frame)
        normalized_depth = resized_depth_frame / max_depth
        depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(normalized_depth, alpha=255), cv2.COLORMAP_JET)
        cv2.imshow("Depth Frame", depth_colormap)

        # 保存图像到文件
        if cv2.waitKey(1) == ord('c'):
            count += 1
            image_filename = f"captured_image_{count}.jpg"
            cv2.imwrite(image_filename, resized_frame)
            print("Image saved to 'captured_image.jpg'")

        # 按 'q' 键退出
        if cv2.waitKey(1) == ord('q'):
            break

cv2.destroyAllWindows()